from binary_agent import binary_agent
from packetholder import packetholder
from constants import PCMD, PSIZE, MAGIC, SYSTEM

class qos_gate(binary_agent):
  def __init__(self, mbase=0, dbase=0):
    super(qos_gate, self).__init__()

    self.m_credits = SYSTEM.QOS_MCREDIT_I
    self.d_credits = SYSTEM.QOS_DCREDIT_I

    self.m_base = mbase
    self.d_base = dbase

    self.net_util_bw = 0.0

  def get_router_util_bw(self):
    return self.output.stat_all_util_bw

  def set_net_util_bw(self, net_util_bw):
    self.net_util_bw = net_util_bw

  def get_net_util_bw(self):
    return self.net_util_bw

  def qos_admit_ratio(self):
    router_admit = 100.0 - 100.0 * self.get_router_util_bw()
    net_admit = 100.0 - 100.0 * self.get_net_util_bw()

    qos_admit = router_admit * (net_admit**2 + router_admit**2) / 2000000.0
    return qos_admit

  def replenish(self):
    # Replenish credits
    r = self.qos_admit_ratio()
    self.m_credits += self.m_base + SYSTEM.QOS_MCREDIT_R * r
    self.d_credits += self.d_base + SYSTEM.QOS_DCREDIT_R * r

  def handle_packet(self, pkt):
    cost = 0
    dma_cost = 0

    # only admission control requests
    if(pkt.cmd in [PCMD.REQ_READ, PCMD.REQ_EXCL]):
      cost = 1
      if(self.m_credits < cost):
        return False

    elif(pkt.cmd in [PCMD.DMA_READ, PCMD.DMA_WRITE]):
      dma_cost = 1
      if(self.d_credits < dma_cost):
        return False

    if(self.out_holders[pkt.priority].accept(pkt, MAGIC.MAX_PSIZE)):
      self.m_credits -= cost
      self.d_credits -= dma_cost
      return True
    else:
      return False
